#!/usr/bin/env python3


import numpy as np


def contained_in(x, a, b):
    return (a<=x) * (x<b)


def dgm(x):
    # Daubechies-Grossmann-Meyer's helper
    return (x>=1) + contained_in(x, 0, 1) * np.sin(x*np.pi/2)**2


def Fpsi(w, nu=dgm):
    # Fourier Transform of Meyer wavelet
    w = np.abs(w)
    return contained_in(w,1/8,1/4) * np.sin(np.pi/2 *nu(8*w-1)) + contained_in(w, 1/4, 1/2) * np.cos(np.pi/2 * nu(4*w-1))


from mymath.fft import *
from ft import *

def wc(f, Fpsi, krange=(-10,10), m=0, Nw=10000):
    # wavelet wrt m
    # assert Nw is large enough
    a = 2**(-m-1)
    Ff, ws = cft(f, trange=(0,1), Nt=None, wrange=(-a, a), Nw=Nw)
    ws = np.linspace(-0.5,0.5, Nw)
    return 2**(-m/2)*fc(Ff[::-1] * Fpsi(ws), krange=krange, tstart='s')


m = 2
t = np.linspace(0,1,1000)
f = np.sin(2*t)*2
n_range=(-20, 20)
k = np.arange(*n_range)
# wc = np.imag(wc(f, Fpsi, krange=n_range, m=m))
# import matplotlib.pyplot as plt
# fig = plt.figure()
# ax = fig.add_subplot(111)
# ax.stem(k, wc)
# plt.show()


def draw_wavelet():
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.subplots(2)
    ws = np.linspace(-.5, .5, 10001)
    F = Fpsi(ws)
    ax[0].plot(ws, F)
    ax[0].set_title('$\hat{\psi}$')
    psi,ts = cft(F)
    ax[1].plot(ts, np.real(psi))
    ax[1].set_title('$\psi$')
    plt.show()

draw_wavelet()

    